import os
from qagen.utils import *
from random import *
import bpy
import mathutils
from mathutils import Euler, Matrix
import math

cate_list = read_category_list()
obj_list = read_obj_info()


def get_object_world_bbox(obj):
    mesh = obj.data
    world_verts = [obj.matrix_world @ v.co for v in mesh.vertices]
    return world_verts

def getBbox(parent_obj):
    all_points = []
    bpy.context.view_layer.update()
    if 'TableDining' in parent_obj.name:
        mesh = parent_obj.data
        worlds_verts = [Euler(parent_obj.rotation_euler, parent_obj.rotation_mode).to_matrix().to_4x4() @ v.co for v in mesh.vertices]
        all_points.extend(worlds_verts)
    for obj in parent_obj.children_recursive:
        if obj.type == 'MESH' and not obj.hide_viewport and not obj.hide_render:
            part_bbox = get_object_world_bbox(obj)
            all_points.extend(part_bbox)
    min_corner = mathutils.Vector((
        min(p.x for p in all_points),
        min(p.y for p in all_points),
        min(p.z for p in all_points)
    ))
    max_corner = mathutils.Vector((
        max(p.x for p in all_points),
        max(p.y for p in all_points),
        max(p.z for p in all_points)
    ))

    return min_corner, max_corner

def fixLocation(obj, current_center, target_center):
    # global depsgraph
    for obj in obj.children_recursive:
        if obj.children == ():
            obj.location = obj.location - current_center + target_center
            # depsgraph = bpy.context.evaluated_depsgraph_get()

def random_position(sz, tablemin, tablemax):
    tx_min = tablemin.x + sz.x / 2
    tx_max = tablemax.x - sz.x / 2
    ty_min = tablemin.y + sz.y / 2
    ty_max = tablemax.y - sz.y / 2
    tx = uniform(tx_min, tx_max)
    ty = uniform(ty_min, ty_max)
    tz = tablemax.z
    return (tx, ty, tz)

def checkOverlap(obj_nbbox):
    n = len(obj_nbbox)
    for i in range(n):
        (a_min, a_max) = obj_nbbox[i]
        axmin, aymin, _ = a_min
        axmax, aymax, _ = a_max
        for j in range(i + 1, n):
            (b_min, b_max) = obj_nbbox[j]
            bxmin, bymin, _ = b_min
            bxmax, bymax, _ = b_max
            if axmin >= bxmax or axmax <= bxmin or aymin >= bymax or aymax <= bymin:
                continue
            return True
    return False

def recursive_select(obj):
    obj.select_set(True)
    for child in obj.children_recursive:
        child.select_set(True)

def delete_sketchfab_models(blend_file_path):
    bpy.ops.wm.open_mainfile(filepath=blend_file_path)
    bpy.ops.object.select_all(action='DESELECT')
    for obj in bpy.data.objects:
        if "Sketchfab_model" in obj.name:
            recursive_select(obj)
    bpy.ops.object.delete()
    bpy.ops.wm.save_mainfile(filepath=blend_file_path)

def random_camera(scene_path):
    bpy.ops.wm.open_mainfile(filepath=scene_path)

    depsgraph = bpy.context.evaluated_depsgraph_get()

    def get_object_world_bbox(obj, depsgraph):
        obj_eval = obj.evaluated_get(depsgraph)
        return [obj_eval.matrix_world @ mathutils.Vector(corner) for corner in obj_eval.bound_box]
    def getBbox(parent_obj, depsgraph):
        all_points = []
        all_points.extend(get_object_world_bbox(parent_obj, depsgraph))
        for child in parent_obj.children_recursive:
            all_points.extend(get_object_world_bbox(child, depsgraph))
        min_corner = mathutils.Vector((
            min(p.x for p in all_points),
            min(p.y for p in all_points),
            min(p.z for p in all_points)
        ))
        max_corner = mathutils.Vector((
            max(p.x for p in all_points),
            max(p.y for p in all_points),
            max(p.z for p in all_points)
        ))
        return min_corner, max_corner

    table_obj = None
    for obj in bpy.data.objects:
        if "TableDining" in obj.name:
            table_obj = obj
            break
    if table_obj is None:
        print("未找到名称包含 'TableDining' 的物体！")
    else:
        print(table_obj.location)
        min_corner, max_corner = getBbox(table_obj, depsgraph)
        print(min_corner, max_corner)
        center = (min_corner + max_corner) / 2.0 + table_obj.location
        width = max_corner.x - min_corner.x
        depth = max_corner.y - min_corner.y
        top_z = max_corner.z

        camera_obj = None
        camrig = None
        for obj in bpy.data.objects:
            if obj.type == 'CAMERA':
                camera_obj = obj
                print(obj.name, obj.type, obj.data.angle_x, obj.data.angle_y)
            if obj.name == 'camrig.0':
                camrig = obj
                print(obj.location, obj.rotation_euler)

        cam_data = camera_obj.data

        fov_x = cam_data.angle_x
        fov_y = cam_data.angle_y

        req_distance_x = (width / 2.0) / math.tan(fov_x / 2.0)
        req_distance_y = (depth / 2.0) / math.tan(fov_y / 2.0)
        req_distance = max(req_distance_x, req_distance_y)
        print("计算所需的距离: %.3f" % req_distance)

        rand_height = uniform(0.0, 0.1)
        rand_ang = randint(45, 60)
        camrig.location = (center.x + (req_distance + rand_height) * math.sin(math.radians(rand_ang)), center.y, top_z + req_distance + rand_height)
        camrig.rotation_euler = (math.radians(rand_ang), 0, math.radians(90))
        print("摄像机新位置:", camrig.location)
        print("摄像机新旋转 (度):", 
              (math.degrees(camrig.rotation_euler.x), 
               math.degrees(camrig.rotation_euler.y), 
               math.degrees(camrig.rotation_euler.z)))

        bpy.ops.wm.save_as_mainfile(filepath=scene_path)

def getglblen(scl):
    scl = max(scl, 15)
    scl = min(scl, 40)
    scl = scl * 0.01 * uniform(0.8, 1.2)
    return scl

def genqa(id):

    scene_id = random_scene()
    scene_id = 'test' #########################
    scene_path = get_scene_path(scene_id)

    bpy.ops.wm.open_mainfile(filepath=scene_path)

    random_camera(scene_path)

    glb_info = []
    obj_num = 3
    for _ in range(obj_num):
        # random in obj_list.keys()
        obj_id = choice(list(obj_list.keys()))
        print(obj_id)
        glb_info.append(obj_list[obj_id])

    imported_objects = []
    for obj in glb_info:
        print(f'importing {obj["path"]}')
        # input()
        bpy.ops.import_scene.gltf(filepath=f'myobj/{obj["path"]}')
        imported_objects.extend(bpy.context.selected_objects)

    bpy.ops.object.mode_set(mode='OBJECT')

    # depsgraph = bpy.context.evaluated_depsgraph_get()

    if not imported_objects:
        raise Exception("import nothing")

    table = None
    for obj in bpy.data.objects:
        if 'TableDining' in obj.name:
            table = obj
            break
    tablemin, tablemax = getBbox(table)
    tablemin += table.location
    tablemax += table.location
    tablesz = tablemax - tablemin
    tablemax.x -= tablesz.x * 0.1

    anc_obj = []
    info_idx = 0
    for obj in imported_objects:
        if obj.parent is None:
            omin, omax = getBbox(obj)
            anc_obj.append(obj)
            omin, omax = getBbox(obj)
            maxlen = getglblen(glb_info[info_idx]['scale'])
            obj.scale *= maxlen / max(omax - omin)
            bpy.context.view_layer.update()
            # depsgraph = bpy.context.evaluated_depsgraph_get()
            omin, omax = getBbox(obj)
            info_idx += 1

    bpy.context.view_layer.update()
    # depsgraph = bpy.context.evaluated_depsgraph_get()

    obj_bbox = []
    for obj in anc_obj:
        omin, omax = getBbox(obj)
        obj_bbox.append(omax - omin)

    while True:
        obj_nbbox = []
        for obj, sz in zip(anc_obj, obj_bbox):
            target_location = random_position(sz, tablemin, tablemax)
            omin, omax = getBbox(obj)
            current_center = mathutils.Vector(((omax.x + omin.x) / 2, (omax.y + omin.y) / 2, omin.z))
            target_center = mathutils.Vector(target_location)
            obj.location = obj.location - current_center + target_center
            bpy.context.view_layer.update()
            # depsgraph = bpy.context.evaluated_depsgraph_get()
            omin, omax = getBbox(obj)
            obj_nbbox.append((omin, omax))
        if not checkOverlap(obj_nbbox):
            break

    # bpy.ops.wm.save_as_mainfile(filepath=...)
    bpy.ops.wm.save_as_mainfile()

    render(scene_id)

    delete_sketchfab_models(scene_path)
    



if __name__ == "__main__":
    num = 1

    for _ in range(num):
        id = random_hex()
        os.makedirs(f'data/{id}', exist_ok=True)
        genqa(id)
